// 定义脚本，把预置模型安装到本地机器，具体包括
// 1. （下载）拷贝模型文件到本地临时目录中（和labelstudio.db同一级目录）
// 2. 更新models_cloud表数据，方便本地可以加载
const os = require('os');
const fs = require('fs');
const path = require('path');
const sqlite3 = require('sqlite3').verbose();
const axios = require('axios');
const yargs = require('yargs');
const archiver = require('archiver');
const sdk = require('@baiducloud/sdk');
const BosClient = sdk.BosClient;

const { ModelsCloudTable } = require('../../main/dbs/tables/tableModelsCloud');
const model_infos = require('./install.json');
const userHomeDir = os.homedir();

let appName = 'labelapp';
let userDataPath;
switch (process.platform) {
    case 'darwin': // macOS
        userDataPath = path.join(userHomeDir, 'Library', 'Application Support', appName);
        break;
    case 'win32': // Windows
        userDataPath = path.join(process.env.APPDATA, appName);
        break;
    default:
        userDataPath = path.join(userHomeDir, `.config`, `${appName}`); // Linux and others
}

function initDB() {
    const dbFilePath = path.join(userDataPath, 'models_cloud', 'models_cloud.db');
    const db = new sqlite3.Database(dbFilePath);
    console.log(`## Open ModelsCloudDB OK [${dbFilePath}]`);
    return db;
}

async function closeDB(db) {
    return new Promise((resolve, reject) => {
        db.close((err) => {
            if (err) {
                console.error('## Close ModelsCloudDB Failed, Error: ', err);
                reject(err);
            } else {
                console.log('## Close ModelsCloudDB OK!');
                resolve();
            }
        });
    });
}

// 从install.json中读取安装信息
// 并且补全这些def_path, conf_path, script_infer, script_train, script_eval, script_env,
async function insertModelsCloud(models_table) {
    for (let i = 0; i < model_infos.length; i++) {
        let model_info = model_infos[i];
        const model_dir = model_info.dir;
        model_info['def_path'] = path.join('models_cloud', model_dir, 'model_infos');
        model_info['conf_path'] = path.join('models_cloud', model_dir, 'config.json');
        model_info['script_infer'] = path.join('models_cloud', model_dir, 'infer.py');
        model_info['script_train'] = path.join('models_cloud', model_dir, 'train.py');
        model_info['script_jsproxy'] = path.join('models_cloud', model_dir, 'proxy.js');
        model_info['script_eval'] = path.join('models_cloud', model_dir, 'eval.py');
        model_info['script_env'] = path.join('models_cloud', model_dir, 'env.py');
        console.log(`Install Model [${model_info.name}]`);
        await models_table.insert(model_info);
    }
}

async function clearModelsCloud(models_table) {
    await models_table.clearModelsCloud();
}

function deleteDirectorySync(dirPath) {
    if (fs.existsSync(dirPath)) {
        fs.readdirSync(dirPath).forEach((file, index) => {
            const currentPath = path.join(dirPath, file);
            if (fs.lstatSync(currentPath).isDirectory()) {
                deleteDirectorySync(currentPath);
            } else {
                fs.unlinkSync(currentPath);
            }
        });
        fs.rmdirSync(dirPath);
    }
}

// 拷贝模型相关的代码到本地目录中（开发阶段通过脚本执行，正式产品需要从云端下载打包的）
function copyDirectorySync(src, dest, blacklist) {
    // 先删除目标目录（如果存在）
    deleteDirectorySync(dest);

    // 重新创建目标目录
    fs.mkdirSync(dest, { recursive: true });

    let entries = fs.readdirSync(src, { withFileTypes: true });
    entries.forEach((entry) => {
        if (blacklist.includes(entry.name)) {
            console.log(`跳过黑名单中的项：${entry.name}`);
            return;
        }

        let srcPath = path.join(src, entry.name);
        let destPath = path.join(dest, entry.name);

        if (entry.isDirectory()) {
            copyDirectorySync(srcPath, destPath, blacklist);
        } else {
            fs.copyFileSync(srcPath, destPath);
        }
    });
}

async function installModelsCloud() {
    let src_models_dir = path.join(__dirname, '..', '..', 'models');
    console.log(`src_models_dir = `, src_models_dir);
    let dst_models_dir = path.join(userDataPath, 'models_cloud', 'models');
    console.log(`dst_models_dir = `, dst_models_dir);
    copyDirectorySync(src_models_dir, dst_models_dir, ['work_dirs', '__pycache__', 'output', 'scripts']);
}

// BOS相关的处理函数
const config = {
    endpoint: 'http://bd.bcebos.com',
    credentials: {
        ak: '6356708ce8d342bf965f2047d7c3cdb2',
        sk: 'd4837f696fdf4cf4abf11a80a1067abb'
    }
};

// BOS相关操作
const bucket = 'label-anything';
const client = new BosClient(config);

async function getCurrentVersion(indexFilePath) {
    try {
        let index_url = 'https://label-anything.bd.bcebos.com/' + indexFilePath;
        let resp = await axios.get(index_url);
        return resp.data.trim();
    } catch (error) {
        // console.log(error);
        return null;
    }
}

function incrementVersion(version) {
    const parts = version.split('.').map(Number);
    parts[2] += 1;
    return parts.join('.');
}

async function uploadFileToBOS(localFilePath, remoteFileName) {
    try {
        const response = await client.putObjectFromFile(bucket, remoteFileName, localFilePath);
        // console.log(`${remoteFileName} uploaded successfully`);
        return response;
    } catch (error) {
        console.error(`Failed to upload ${remoteFileName}:`, error);
        throw error;
    }
}

// 压缩对应的目录
async function zipDirectory(sourceDir, outPath) {
    const archive = archiver('zip', { zlib: { level: 9 } });
    const stream = fs.createWriteStream(outPath);

    return new Promise((resolve, reject) => {
        archive
            .directory(sourceDir, false)
            // .directory(sourceDir, path.basename(sourceDir))
            .on('error', (err) => reject(err))
            .pipe(stream);

        stream.on('close', () => resolve());
        archive.finalize();
    });
}

async function zipFilesAndFolders(inputList, outPath) {
    const archive = archiver('zip', { zlib: { level: 9 } });
    const stream = fs.createWriteStream(outPath);

    return new Promise((resolve, reject) => {
        archive.on('error', (err) => reject(err)).pipe(stream);

        for (const item of inputList) {
            const itemPath = path.resolve(item);

            // 判断是文件还是文件夹
            if (fs.statSync(itemPath).isDirectory()) {
                archive.directory(itemPath, path.basename(itemPath));
            } else {
                archive.file(itemPath, { name: path.basename(itemPath) });
            }
        }

        stream.on('close', () => resolve());
        archive.finalize();
    });
}

async function main() {
    const argv = yargs
        .option('mode', {
            alias: 'm',
            description: 'Specify the operation mode: local or remote',
            choices: ['local', 'remote'],
            default: 'local'
        })
        .help()
        .alias('help', 'h').argv;

    const mode = argv.mode;
    console.log(`## Deploying to ${mode}`);

    // 拷贝模型的配置和训练脚本
    await installModelsCloud();

    // 构建数据库信息
    try {
        let db = initDB();
        let models_table = new ModelsCloudTable(null, db);
        await models_table.init();
        await clearModelsCloud(models_table);
        await insertModelsCloud(models_table);
        await closeDB(db);
    } catch (err) {
        console.log(`## Error:`, err);
    }

    if (mode == 'remote') {
        // 获取当前版本号
        const indexFilePath = '/labelapp/models_cloud/VERSION';
        let currentVersion = await getCurrentVersion(indexFilePath);
        console.log(`## Current version: ${currentVersion}`);
        if (currentVersion == null) {
            console.log('## Current version not found, please check network connection.');
            return;
        }

        // 增量版本号，并更新到dist输出目录下
        const newVersion = incrementVersion(currentVersion);
        const newVersionPath = path.join(__dirname, '..', 'VERSION');
        const installVersionPath = path.join(userDataPath, 'models_cloud', 'VERSION');
        console.log(`## New version: ${newVersion}`);
        fs.writeFileSync(newVersionPath, newVersion);
        fs.writeFileSync(installVersionPath, newVersion);

        // 压缩文件并使用新版本号命名
        const installDir = path.join(userDataPath, 'models_cloud', 'models');
        const dbFilePath = path.join(userDataPath, 'models_cloud', 'models_cloud.db');
        const zipFilePath = path.join(__dirname, `models-${newVersion}.zip`);
        // await zipDirectory(installDir, zipFilePath);
        const zipList = [installDir, dbFilePath];
        await zipFilesAndFolders(zipList, zipFilePath);
        console.log(`## Directory compressed to ${zipFilePath}`);

        // 上传压缩文件到百度云BOS
        let remotePath = `/labelapp/models_cloud/models-${newVersion}.zip`;
        await uploadFileToBOS(zipFilePath, remotePath);
        console.log(`## File ${remotePath} uploaded to BOS.`);

        // 上传新的版本号到VERSION
        await uploadFileToBOS(newVersionPath, indexFilePath);
        console.log(`## Updated version info to BOS.`);

        // 清理临时压缩文件
        fs.unlinkSync(zipFilePath);
        console.log('## Temporary file cleaned up.');
    }
}

main();
